import os
from pathlib import Path
import numpy as np
from ultralytics.utils.ops import segments2boxes as ultralytics_segments2boxes

# 请确保在实际使用时该函数可用
def segments2boxes(segments):
    processed_segments = []
    for class_index, coords in segments:
        # 跳过类别索引，将坐标转换为 numpy 数组并重塑为 (n, 2) 形状
        coord_array = np.array(coords).reshape(-1, 2)
        processed_segments.append(coord_array)
    # 调用 ultralytics 的 segments2boxes 函数
    return ultralytics_segments2boxes(processed_segments)

def process_txt_files(input_dir, output_dir):
    input_path = Path(input_dir)
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    for txt_file in input_path.glob('*.txt'):
        # STEP1：读取segments list
        segments = []
        with open(txt_file, 'r') as f:
            for line in f:   # 因为不同的segment会分为不同的line，所以需要逐行读取
                parts = list(map(float, line.strip().split()))
                class_index = parts[0]
                coords = parts[1:]
                segments.append((class_index, coords))
        print(f'文件 {txt_file} 的 segments 数量为: {len(segments)}')

        # STEP2：将segments转化为boxs
        box_lines = []
        boxes = segments2boxes(segments)
        for i, box in enumerate(boxes):
            class_index = int(segments[i][0])   #类别是int
            box_str = ' '.join(map(str, [class_index] + box.tolist()))
            box_lines.append(box_str)
            print(box_str)  # 打印每个 box_lines每个 box_lines

        # STEP3：将boxs写入新的txt文件
        output_file = output_path / txt_file.name
        with open(output_file, 'w') as f:
            f.write('\n'.join(box_lines))


if __name__ == '__main__':
    input_directory = 'YOLO/SAM/images_auto_annotate_labels'
    output_directory = 'YOLO/SAM/images_auto_annotate_boxs'
    process_txt_files(input_directory, output_directory)